import pytesseract
from PIL import Image
import cv2
import numpy as np
import re
import os
from collections import Counter


class UniversalChineseOCR:
    def __init__(self):
        # 通用中文OCR配置，不使用字符白名单
        self.configs = [
            r'--oem 1 --psm 6 -c preserve_interword_spaces=1 -l chi_sim+eng',  # 保留空格
            r'--oem 1 --psm 7 -c preserve_interword_spaces=1 -l chi_sim+eng',  # 单行文本
            r'--oem 3 --psm 6 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789()（）[]【】{}《》\u4e00-\u9fff -l chi_sim+eng',
            # 字符白名单
            r'--oem 1 --psm 4 -c preserve_interword_spaces=1 -l chi_sim+eng',  # 多列文本
        ]

    def preprocess_image_universal(self, image_path):
        """改进的图像预处理，增强括号识别"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像文件: {image_path}")

        # 1. 自动调整图像大小 - 确保足够大的分辨率来识别细节
        height, width = img.shape[:2]
        if max(height, width) > 2000:  # 如果图像太大，缩小
            scale = 1500 / max(height, width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
        elif max(height, width) < 1000:  # 如果图像太小，放大到更大尺寸
            scale = 1800 / max(height, width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_CUBIC)

        # 2. 多种预处理方法
        processed_images = []

        # 方法1: 标准灰度 + Otsu二值化 + 形态学操作
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        _, thresh1 = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        # 添加形态学操作来清理噪声，但保留细小结构
        kernel = np.ones((1, 1), np.uint8)
        thresh1 = cv2.morphologyEx(thresh1, cv2.MORPH_CLOSE, kernel)
        processed_images.append(('otsu_morph', thresh1))

        # 方法2: 自适应阈值 - 更适合处理括号等细小结构
        thresh2 = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                        cv2.THRESH_BINARY, 11, 2)
        processed_images.append(('adaptive_tuned', thresh2))

        # 方法3: 对比度增强 + 锐化 - 增强括号边缘
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)

        # 锐化操作增强边缘细节
        kernel_sharpen = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
        sharpened = cv2.filter2D(enhanced, -1, kernel_sharpen)

        _, thresh3 = cv2.threshold(sharpened, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('enhanced_sharpen', thresh3))

        # 方法4: 去噪 + 直方图均衡化 - 保持细节
        denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)
        equalized = cv2.equalizeHist(denoised)
        _, thresh4 = cv2.threshold(equalized, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('equalized_denoised', thresh4))

        # 方法5: 专门针对细小字符的处理
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        _, thresh5 = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('fine_detail', thresh5))

        return processed_images

    def extract_text_multiple_methods(self, image_path):
        """使用多种方法提取文本"""
        processed_images = self.preprocess_image_universal(image_path)

        all_results = []

        for method_name, img in processed_images:
            for config in self.configs:
                try:
                    # 转换为PIL图像进行OCR
                    pil_img = Image.fromarray(img)
                    text = pytesseract.image_to_string(pil_img, config=config)

                    if text.strip():
                        all_results.append({
                            'method': f"{method_name}_{config.split('psm ')[1][0]}",
                            'text': text,
                            'char_count': len(text.strip()),
                            'chinese_char_count': len(re.findall(r'[\u4e00-\u9fff]', text))
                        })

                except Exception as e:
                    continue

        return all_results

    def select_best_result(self, results):
        """选择最佳的识别结果"""
        if not results:
            return None

        # 根据中文字符数量和总字符数量评分
        scored_results = []
        for result in results:
            score = result['chinese_char_count'] * 2 + result['char_count']
            # 如果包含常见关键词，加分
            keywords = ['根据', '研究', '实现', '无法', '选项', '答案']
            keyword_bonus = sum(10 for keyword in keywords if keyword in result['text'])
            score += keyword_bonus

            scored_results.append((score, result))

        # 选择分数最高的结果
        scored_results.sort(key=lambda x: x[0], reverse=True)
        return scored_results[0][1] if scored_results else None

    def extract_question_blocks(self, text):
        """改进的题目块提取"""
        lines = [line.strip() for line in text.split('\n') if line.strip()]

        # 寻找可能的题目开始行
        question_starts = []
        for i, line in enumerate(lines):
            # 题目通常包含这些特征
            if (len(line) > 5 and
                    any(char in line for char in ['根据', '研究', '问题', '题目', '？', '?']) or
                    re.search(r'^[A-D][\.\、\．:]', line) or  # 包含选项标识符
                    re.search(r'[0-9]{1,3}/[0-9]{1,3}', line)):  # 包含页码
                question_starts.append(i)

        # 提取题目块
        question_blocks = []
        for start_idx in question_starts[:5]:  # 增加检查的数量
            block_lines = []
            option_count = 0

            for i in range(start_idx, min(start_idx + 20, len(lines))):  # 增加行数范围
                line = lines[i]
                block_lines.append(line)

                # 统计选项数量 - 改进正则表达式
                if re.match(r'^[A-D][\.\、\．:\s]', line.strip()):
                    option_count += 1

                # 如果已经收集了4个选项或者遇到答案信息，停止
                if option_count >= 4 or any(keyword in line for keyword in ['正确答案', '答案', '解析']):
                    # 继续读取几行以防答案在同一区域
                    for j in range(i + 1, min(i + 4, len(lines))):
                        next_line = lines[j]
                        if any(keyword in next_line for keyword in ['正确答案', '答案', '解析']):
                            block_lines.append(next_line)
                        elif re.match(r'^[A-D][\.\、\．:\s]', next_line.strip()):
                            block_lines.append(next_line)
                    break

            # 更宽松的有效题目块判断条件
            if len(block_lines) >= 2 and (option_count >= 2 or
                                          any('答案' in line for line in block_lines)):
                question_blocks.append({
                    'start_line': start_idx,
                    'lines': block_lines,
                    'option_count': option_count,
                    'text': '\n'.join(block_lines)
                })

        return question_blocks

    def analyze_text_structure(self, text):
        """分析文本结构"""
        lines = [line.strip() for line in text.split('\n') if line.strip()]

        structure = {
            'header': [],
            'content_blocks': [],
            'questions': [],
            'options_blocks': [],
            'footer': []
        }

        current_block = []
        current_type = 'content'

        for i, line in enumerate(lines):
            # 检测各部分类型
            if i == 0 and (re.search(r'\d+/\d+', line) or '案例' in line or '题目' in line):
                structure['header'].append(line)
                continue

            # 检测选项块
            if re.match(r'^[A-D][\.\、\．\s]', line):
                if current_block and current_type == 'content':
                    structure['content_blocks'].append('\n'.join(current_block))
                    current_block = []
                current_type = 'options'
                structure['options_blocks'].append(line)
                continue

            # 检测问题行
            if '根据' in line and '研究' in line and len(line) > 10:
                if current_block:
                    structure['content_blocks'].append('\n'.join(current_block))
                current_block = [line]
                current_type = 'question'
                continue

            # 检测答案信息
            if any(keyword in line for keyword in ['正确答案', '我的答案', '参考解析']):
                structure['footer'].append(line)
                continue

            # 添加到当前块
            current_block.append(line)

        # 添加最后一个块
        if current_block:
            if current_type == 'question':
                structure['questions'].append('\n'.join(current_block))
            else:
                structure['content_blocks'].append('\n'.join(current_block))

        return structure

    def process_single_image(self, image_path):
        """处理单张图片"""
        print(f"处理图片: {os.path.basename(image_path)}")

        # 提取文本
        results = self.extract_text_multiple_methods(image_path)

        if not results:
            print("  ⚠️  未能识别到文本")
            return None

        # 选择最佳结果
        best_result = self.select_best_result(results)

        if not best_result:
            print("  ⚠️  无法选择最佳结果")
            return None

        print(f"  ✅ 使用方法: {best_result['method']}")
        print(f"  📊 字符数: {best_result['char_count']}, 中文字符: {best_result['chinese_char_count']}")

        # 分析文本结构
        structure = self.analyze_text_structure(best_result['text'])

        return {
            'file_name': os.path.basename(image_path),
            'method': best_result['method'],
            'full_text': best_result['text'],
            'structure': structure,
            'question_blocks': self.extract_question_blocks(best_result['text'])
        }

    def batch_process_images(self, image_folder):
        """批量处理图片文件夹"""
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']

        # 获取所有图片文件
        image_files = []
        for file in os.listdir(image_folder):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(image_folder, file))

        print(f"找到 {len(image_files)} 张图片")

        results = []
        for image_path in image_files:
            result = self.process_single_image(image_path)
            if result:
                results.append(result)
            print("-" * 50)

        return results

    def save_results(self, results, output_dir="ocr_results"):
        """保存识别结果"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 保存每个文件的详细结果
        for result in results:
            base_name = os.path.splitext(result['file_name'])[0]
            output_file = os.path.join(output_dir, f"{base_name}_result.txt")

            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(f"文件: {result['file_name']}\n")
                f.write(f"识别方法: {result['method']}\n")
                f.write("=" * 60 + "\n\n")

                f.write("完整识别文本:\n")
                f.write("-" * 40 + "\n")
                f.write(result['full_text'] + "\n\n")

                f.write("文本结构分析:\n")
                f.write("-" * 40 + "\n")
                structure = result['structure']
                if structure['header']:
                    f.write("头部信息:\n")
                    for line in structure['header']:
                        f.write(f"  {line}\n")
                    f.write("\n")

                if structure['content_blocks']:
                    f.write("内容块:\n")
                    for i, block in enumerate(structure['content_blocks']):
                        f.write(f"块 {i + 1}:\n{block}\n\n")

                if structure['questions']:
                    f.write("问题部分:\n")
                    for i, question in enumerate(structure['questions']):
                        f.write(f"问题 {i + 1}:\n{question}\n\n")

                if structure['options_blocks']:
                    f.write("选项:\n")
                    for option in structure['options_blocks']:
                        f.write(f"  {option}\n")
                    f.write("\n")

                if structure['footer']:
                    f.write("底部信息:\n")
                    for line in structure['footer']:
                        f.write(f"  {line}\n")

                if result['question_blocks']:
                    f.write("\n提取的题目块:\n")
                    f.write("-" * 40 + "\n")
                    for i, block in enumerate(result['question_blocks']):
                        f.write(f"题目块 {i + 1} (选项数: {block['option_count']}):\n")
                        f.write(block['text'] + "\n\n")

        # 保存汇总报告
        summary_file = os.path.join(output_dir, "summary_report.txt")
        with open(summary_file, 'w', encoding='utf-8') as f:
            f.write("OCR识别汇总报告\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"处理图片数量: {len(results)}\n")
            f.write(f"成功识别数量: {len(results)}\n\n")

            for result in results:
                f.write(f"文件: {result['file_name']}\n")
                f.write(f"方法: {result['method']}\n")
                f.write(f"字符数: {len(result['full_text'])}\n")
                f.write(f"题目块数: {len(result['question_blocks'])}\n")
                f.write("-" * 30 + "\n")

        print(f"\n结果已保存到目录: {output_dir}")
        print(f"汇总报告: {summary_file}")


def main():
    ocr = UniversalChineseOCR()

    # 选择处理模式
    print("请选择处理模式:")
    print("1. 处理单张图片")
    print("2. 批量处理图片文件夹")

    choice = input("请输入选择 (1 或 2): ").strip()

    if choice == "1":
        image_path = input("请输入图片路径: ").strip()
        if not os.path.exists(image_path):
            print("文件不存在!")
            return

        result = ocr.process_single_image(image_path)
        if result:
            print("\n" + "🎯" * 50)
            print("识别结果")
            print("🎯" * 50)
            print(f"文件: {result['file_name']}")
            print(f"方法: {result['method']}")
            print("\n完整文本:")
            print("-" * 60)
            print(result['full_text'])
            print("-" * 60)

            # 保存单个结果
            ocr.save_results([result])

    elif choice == "2":
        folder_path = input("请输入图片文件夹路径: ").strip()
        if not os.path.exists(folder_path):
            print("文件夹不存在!")
            return

        results = ocr.batch_process_images(folder_path)
        if results:
            ocr.save_results(results)
        else:
            print("没有成功识别任何图片")

    else:
        print("无效选择!")


if __name__ == "__main__":
    main()